﻿using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using Newtonsoft.Json;
using System;
using System.Collections.Generic;
using System.Drawing;
using System.Drawing.Imaging;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ConsoleApp
{
    public class YoloOnnxRuntimeDetector
    {
        // 存储类别名称的字典
        private readonly Dictionary<int, string> classNames;
        // ONNX推理会话
        private readonly InferenceSession session;
        // 输入图像的宽度和高度
        private static readonly int INPUT_WIDTH = 640;
        private static readonly int INPUT_HEIGHT = 640;

        public YoloOnnxRuntimeDetector(string modelPath, string classesYamlPath)
        {
            this.classNames = FlygoDesktop.Utility.ProfilesHelper.LoadClassNamesFromYaml(classesYamlPath);
            // 创建ONNX推理会话
            this.session = new InferenceSession(modelPath);
        }
        
        public string Detect(string imagePath, float confidenceThreshold = 0.5f, float nmsThreshold = 0.45f)
        {
            // 加载并预处理图像
            var (inputTensor, originalWidth, originalHeight) = LoadImage(imagePath);

            // 准备输入参数
            var inputs = new List<NamedOnnxValue> { NamedOnnxValue.CreateFromTensor("images", inputTensor) };

            // 运行模型推理
            using (var results = this.session.Run(inputs))
            {
                // 处理输出结果
                var outputTensor = results.First().AsTensor<float>();
                var predictions = ProcessOutput(outputTensor, originalWidth, originalHeight, this.classNames, confidenceThreshold, nmsThreshold);

                // 绘制检测结果
                // DrawPredictions(imagePath, predictions, "D:/output-1.jpg", this.classNames);
                var finalResult = new
                {
                    results = predictions,
                    error_code = 0
                };
                return Newtonsoft.Json.JsonConvert.SerializeObject(finalResult, Formatting.Indented);
            }
        }
        // 加载并预处理图像（使用LockBits优化性能）
        private static (DenseTensor<float> tensor, int width, int height) LoadImage(string imagePath)
        {
            using (var image = Image.FromFile(imagePath))
            {
                int originalWidth = image.Width;
                int originalHeight = image.Height;

                // 调整图像大小为模型输入尺寸（640x640）
                using (var resizedImage = ResizeImage(image, INPUT_WIDTH, INPUT_HEIGHT))
                using (var bitmap = new Bitmap(resizedImage))
                {
                    // 使用LockBits提高性能
                    var bitmapData = bitmap.LockBits(
                        new Rectangle(0, 0, bitmap.Width, bitmap.Height),
                        ImageLockMode.ReadOnly,
                        PixelFormat.Format24bppRgb);

                    // 创建张量 [1, 3, 640, 640] (批次, 通道, 高, 宽)
                    var tensor = new DenseTensor<float>(new[] { 1, 3, INPUT_HEIGHT, INPUT_WIDTH });

                    try
                    {
                        unsafe
                        {
                            byte* ptr = (byte*)bitmapData.Scan0;
                            int stride = bitmapData.Stride;

                            for (int y = 0; y < bitmap.Height; y++)
                            {
                                byte* row = ptr + (y * stride);

                                for (int x = 0; x < bitmap.Width; x++)
                                {
                                    // 注意：模型需要BGR顺序（OpenCV默认）
                                    // 内存布局: [B, G, R]
                                    byte b = row[x * 3];
                                    byte g = row[x * 3 + 1];
                                    byte r = row[x * 3 + 2];

                                    // 归一化处理 (0-255 -> 0.0-1.0)
                                    tensor[0, 0, y, x] = r / 255.0f; // R通道
                                    tensor[0, 1, y, x] = g / 255.0f; // G通道
                                    tensor[0, 2, y, x] = b / 255.0f; // B通道
                                }
                            }
                        }
                    }
                    finally
                    {
                        bitmap.UnlockBits(bitmapData);
                    }

                    return (tensor, originalWidth, originalHeight);
                }
            }
        }
        // 调整图像大小
        private static Image ResizeImage(Image image, int width, int height)
        {
            var destRect = new Rectangle(0, 0, width, height);
            var destImage = new Bitmap(width, height);

            destImage.SetResolution(image.HorizontalResolution, image.VerticalResolution);

            using (var graphics = Graphics.FromImage(destImage))
            {
                graphics.CompositingMode = System.Drawing.Drawing2D.CompositingMode.SourceCopy;
                graphics.CompositingQuality = System.Drawing.Drawing2D.CompositingQuality.HighQuality;
                graphics.InterpolationMode = System.Drawing.Drawing2D.InterpolationMode.HighQualityBicubic;
                graphics.SmoothingMode = System.Drawing.Drawing2D.SmoothingMode.HighQuality;
                graphics.PixelOffsetMode = System.Drawing.Drawing2D.PixelOffsetMode.HighQuality;

                using (var wrapMode = new System.Drawing.Imaging.ImageAttributes())
                {
                    wrapMode.SetWrapMode(System.Drawing.Drawing2D.WrapMode.TileFlipXY);
                    graphics.DrawImage(image, destRect, 0, 0, image.Width, image.Height, GraphicsUnit.Pixel, wrapMode);
                }
            }

            return destImage;
        }
        // 处理模型输出
        private static List<Prediction> ProcessOutput(Tensor<float> output, int originalWidth, int originalHeight, Dictionary<int, string> classNames, float confidenceThreshold = 0.5f, float nmsThreshold = 0.45f)
        {
            var predictions = new List<Prediction>();

            // 获取输出张量维度 [batch, num_boxes, 5 + num_classes]
            var dimensions = output.Dimensions.ToArray();
            if (dimensions.Length != 3)
            {
                throw new InvalidOperationException("输出张量维度错误");
            }

            int numBoxes = dimensions[1];
            int numClasses = dimensions[2] - 5; // 5 = cx, cy, w, h, confidence

            // 解析模型输出
            for (int i = 0; i < numBoxes; i++)
            {
                float confidence = output[0, i, 4];

                if (confidence >= confidenceThreshold)
                {
                    // 找到概率最高的类别
                    int classId = 0;
                    float maxClassProbs = 0;

                    for (int c = 0; c < numClasses; c++)
                    {
                        float classProbs = output[0, i, 5 + c];
                        if (classProbs > maxClassProbs)
                        {
                            maxClassProbs = classProbs;
                            classId = c;
                        }
                    }

                    float score = confidence * maxClassProbs;

                    if (score >= confidenceThreshold && classId < classNames.Count)
                    {
                        // 获取边界框坐标
                        float cx = output[0, i, 0];
                        float cy = output[0, i, 1];
                        float w = output[0, i, 2];
                        float h = output[0, i, 3];

                        // 转换为原始图像坐标 (左上角+宽高)
                        float scaleX = originalWidth / ((float) INPUT_WIDTH);
                        float scaleY = originalHeight / ((float) INPUT_HEIGHT);

                        // 转换为左上角和右下角坐标
                        float ltx = (cx - w / 2) * scaleX;
                        float lty = (cy - h / 2) * scaleY;
                        float rbx = (cx + w / 2) * scaleX;
                        float rby = (cy + h / 2) * scaleY;

                        float width = w * scaleX;
                        float height = h * scaleY;

                        // 确保坐标在图像范围内
                        ltx = Math.Max(0, Math.Min(originalWidth - 1, ltx));
                        lty = Math.Max(0, Math.Min(originalHeight - 1, lty));
                        width = Math.Max(1, Math.Min(originalWidth - ltx, width));
                        height = Math.Max(1, Math.Min(originalHeight - lty, height));

                        string className = classNames.ContainsKey(classId) ? classNames[classId] : "Unknown";
                        predictions.Add(new Prediction
                        {
                            ClassId = classId,
                            label = className,
                            confidence = confidence,
                            Score = score,
                            BBox = new RectangleF(ltx, lty, width, height),
                            x1 = (cx - w / 2) / ((float)INPUT_WIDTH),
                            y1 = (cy - h / 2) / ((float)INPUT_HEIGHT),
                            x2 = (cx + w / 2) / ((float)INPUT_WIDTH),
                            y2 = (cy + h / 2) / ((float)INPUT_HEIGHT),
                            ltx = ltx,
                            lty = lty,
                            rbx = rbx,
                            rby = rby,
                        });
                    }
                }
            }

            // 非极大值抑制
            return NonMaxSuppression(predictions, nmsThreshold);
        }

        // 非极大值抑制
        private static List<Prediction> NonMaxSuppression(List<Prediction> predictions, float threshold)
        {
            var result = new List<Prediction>();

            // 按置信度降序排序
            predictions = predictions.OrderByDescending(p => p.Score).ToList();

            while (predictions.Count > 0)
            {
                var best = predictions[0];
                result.Add(best);
                predictions.RemoveAt(0);

                // 移除重叠度高的边界框
                for (int i = predictions.Count - 1; i >= 0; i--)
                {
                    if (IoU(best.BBox, predictions[i].BBox) > threshold)
                    {
                        predictions.RemoveAt(i);
                    }
                }
            }

            return result;
        }
        // 计算交并比（修正IoU计算）
        private static float IoU(RectangleF a, RectangleF b)
        {
            float areaA = a.Width * a.Height;
            float areaB = b.Width * b.Height;

            float interX1 = Math.Max(a.Left, b.Left);
            float interY1 = Math.Max(a.Top, b.Top);
            float interX2 = Math.Min(a.Right, b.Right);
            float interY2 = Math.Min(a.Bottom, b.Bottom);

            float interWidth = Math.Max(0, interX2 - interX1);
            float interHeight = Math.Max(0, interY2 - interY1);
            float interArea = interWidth * interHeight;

            return interArea / (areaA + areaB - interArea);
        }

        private static void DrawPredictions(string inputImagePath, List<Prediction> predictions, string outputImagePath, Dictionary<int,string> classNames)
        {
            using (var image = Image.FromFile(inputImagePath))
            using (var graphics = Graphics.FromImage(image))
            {
                // 设置绘图质量
                graphics.SmoothingMode = System.Drawing.Drawing2D.SmoothingMode.AntiAlias;
                graphics.InterpolationMode = System.Drawing.Drawing2D.InterpolationMode.HighQualityBicubic;
                graphics.PixelOffsetMode = System.Drawing.Drawing2D.PixelOffsetMode.HighQuality;

                // 绘制边界框和标签
                foreach (var prediction in predictions)
                {
                    var bbox = prediction.BBox;
                    var label = $"{classNames[prediction.ClassId]}: {prediction.Score:F2}";

                    // 绘制边界框
                    using (var pen = new Pen(Color.FromArgb(255, 255, 0, 0), 2))
                    {
                        graphics.DrawRectangle(pen, bbox.X, bbox.Y, bbox.Width, bbox.Height);
                    }


                    // 绘制标签背景
                    using (var font = new Font("Arial", 10, FontStyle.Bold))
                    using (var brush = new SolidBrush(Color.FromArgb(255, 255, 0, 0)))
                    using (var textBrush = new SolidBrush(Color.White))
                    {
                        var textSize = graphics.MeasureString(label, font);
                        var textBackground = new RectangleF(bbox.X, bbox.Y - textSize.Height, textSize.Width, textSize.Height);

                        graphics.FillRectangle(brush, textBackground);
                        graphics.DrawString(label, font, textBrush, bbox.X, bbox.Y - textSize.Height);
                    }
                }

                // 保存结果图像
                image.Save(outputImagePath, ImageFormat.Jpeg);
            }
        }
    }
}
